package binary_tree.leetcode.medium;

import java.util.HashMap;
import java.util.Map;

/**
 * @author bruin_du
 * @description 单词之和
 * @date 2022/8/18 23:52
 **/
//前缀树加哈希表
public class OfferII066_MapSum {
    TreeNode root;
    Map<String, Integer> map;

    class TreeNode {
        TreeNode[] next;
        int times;
        boolean isEnd;

        public TreeNode(int times) {
            next = new TreeNode[26];
            this.times = times;
        }
    }


    /** Initialize your data structure here. */
    public OfferII066_MapSum() {
        root = new TreeNode(0);
        map = new HashMap<>();
    }

    public void insert(String key, int val) {
        char[] arr = key.toCharArray();
        TreeNode cur = root;
        int oldVal = 0;
        boolean flag = false;
        if (map.containsKey(key)) {
            oldVal = map.get(key);
            flag = true;
        }
        for (char ch : arr) {
            int index = ch - 'a';
            if (cur.next[index] == null) {
                cur.next[index] = new TreeNode(val);
            } else{
                cur.next[index].times += val;
            }
            if (flag) cur.next[index].times -= oldVal;
            cur = cur.next[index];
        }
        cur.isEnd = true;
        map.put(key, val);
    }

    public int sum(String prefix) {
        char[] arr = prefix.toCharArray();
        TreeNode cur = root;
        for (char ch : arr) {
            int index = ch - 'a';
            if (cur.next[index] == null) {
                return 0;
            }
            cur = cur.next[index];
        }
        return cur.times;
    }
}

//前缀哈希映射
class MapSum {
    Map<String, Integer> map;
    Map<String, Integer> prefixmap;

    public MapSum() {
        map = new HashMap<>();
        prefixmap = new HashMap<>();
    }

    public void insert(String key, int val) {
        int delta = val - map.getOrDefault(key, 0);
        map.put(key, val);
        for (int i = 1; i <= key.length(); ++i) {
            String currprefix = key.substring(0, i);
            prefixmap.put(currprefix, prefixmap.getOrDefault(currprefix, 0) + delta);
        }
    }

    public int sum(String prefix) {
        return prefixmap.getOrDefault(prefix, 0);
    }
}